{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c5dbcc9ae0c"
      },
      "source": [
        "##### Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "906e07f6e562"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "52d7c0d0e8fe"
      },
      "source": [
        "*Contributed by the Weaviate team*\n",
        "\n",
        "## Weaviate Query Agent with Gemini API\n",
        "\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/weaviate/query-agent-as-a-tool.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" height=30/></a>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "14846c0e4bf3"
      },
      "source": [
        "This notebook will show you how to define the Weaviate Query Agent as a tool through the Gemini API.\n",
        "\n",
        "### Requirements\n",
        "1. Weaviate Cloud instance (WCD): The Weaviate Query Agent is only accessible through WCD at the moment. You can create a serverless cluster or a free 14-day sandbox [here](https://console.weaviate.cloud/).\n",
        "2. Have a GCP project and Gemini API key (generate one [here](https://aistudio.google.com/))\n",
        "3. Install the Google Gen AI SDK with `pip install --upgrade --quiet google-genai`\n",
        "4. Install the Weaviate Python client and the agents sub-package with `pip install weaviate-client[agents]`\n",
        "5. You'll need a Weaviate cluster with data. If you don't have one, check out [this notebook](integrations/Weaviate-Import-Example.ipynb) to import the Weaviate Blogs.\n",
        "\n",
        "Connect with us and let us know if you have any questions!\n",
        "\n",
        "Erika's accounts:\n",
        "* [Follow on X](https://x.com/ecardenas300)\n",
        "* [Connect on LinkedIn](https://www.linkedin.com/in/erikacardenas300/)\n",
        "\n",
        "Patrick's accounts:\n",
        "* [Follow on X](https://x.com/patloeber)\n",
        "* [Connect on LinkedIn](https://www.linkedin.com/in/patrick-l%C3%B6ber-403022137/)\n",
        "\n",
        "Connor's accounts:\n",
        "* [LinkedIn](https://www.linkedin.com/in/connor-shorten-34923a178/)\n",
        "* [X](https://x.com/CShorten30)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2d2ed01d4df7"
      },
      "source": [
        "### Install libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "57e03c3fd5ef"
      },
      "outputs": [],
      "source": [
        "%pip install -U google-genai\n",
        "%pip install -U \"weaviate-client[agents]\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "55d9f2c341da"
      },
      "source": [
        "### Import libraries and keys"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88bd6e96457a"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import weaviate\n",
        "from weaviate_agents.query import QueryAgent\n",
        "\n",
        "from google import genai\n",
        "from google.genai import types"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9f569fd18f0"
      },
      "source": [
        "### Set you API keys and Weaviate URL"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c6978d8cb04"
      },
      "outputs": [],
      "source": [
        "os.environ[\"WEAVIATE_URL\"] = \"\"\n",
        "os.environ[\"WEAVIATE_API_KEY\"] = \"\"\n",
        "os.environ[\"GOOGLE_API_KEY\"] = \"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "34f0d5c7d990"
      },
      "source": [
        "### Create API client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2911cbd44f2d"
      },
      "outputs": [],
      "source": [
        "client = genai.Client()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "44e7c6bffd55"
      },
      "source": [
        "### Define Query Agent function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f9a410f2f8c5"
      },
      "outputs": [],
      "source": [
        "def query_agent_request(query: str) -> str:\n",
        "    \"\"\"\n",
        "    Send a query to the database and get the response.\n",
        "\n",
        "    Args:\n",
        "        query (str): The question or query to search for in the database. This can be any natural language question related to the content stored in the database.\n",
        "\n",
        "    Returns:\n",
        "        str: The response from the database containing relevant information.\n",
        "    \"\"\"\n",
        "\n",
        "    # connect to your Weaviate Cloud instance\n",
        "    weaviate_client = weaviate.connect_to_weaviate_cloud(\n",
        "        cluster_url=os.getenv(\"WEAVIATE_URL\"),\n",
        "        auth_credentials=weaviate.auth.AuthApiKey(os.getenv(\"WEAVIATE_API_KEY\")),\n",
        "        headers={  # add the API key to the model provider from your Weaviate collection, for example `headers={\"X-Goog-Studio-Api-Key\": os.getenv(\"GEMINI_API_KEY\")}`\n",
        "        }\n",
        "    )\n",
        "\n",
        "    # connect the query agent to your Weaviate collection(s)\n",
        "    query_agent = QueryAgent(\n",
        "        client=weaviate_client,\n",
        "        collections=[\"WeaviateBlogChunks\"]\n",
        "    )\n",
        "    return query_agent.run(query).final_answer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6aadcbfb3d0f"
      },
      "source": [
        "### Configure Tool"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0c55c00e19e4"
      },
      "outputs": [],
      "source": [
        "config = types.GenerateContentConfig(tools=[query_agent_request])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90b697087def"
      },
      "source": [
        "### Query Time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "edf15d21bb0c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.11/site-packages/pydantic/main.py:426: UserWarning: Pydantic serializer warnings:\n",
            "  Expected `enum` but got `str` with value `'STRING'` - serialized value may not be as expected\n",
            "  return self.__pydantic_serializer__.to_python(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "To deploy Weaviate with Docker, you need to:\n",
            "\n",
            "1.  Install Docker and Docker Compose.\n",
            "2.  Obtain the Weaviate Docker image using:\n",
            "    ```bash\n",
            "    docker pull cr.weaviate.io/semitechnologies/weaviate:latest\n",
            "    ```\n",
            "3.  Prepare a `docker-compose.yml` file, which you can generate using the Weaviate configuration tool or example files from the documentation.\n",
            "4.  Start Weaviate using either:\n",
            "    *   Directly with Docker:\n",
            "        ```bash\n",
            "        docker run -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:latest\n",
            "        ```\n",
            "    *   Using Docker Compose:\n",
            "        ```bash\n",
            "        docker-compose up -d\n",
            "        ```\n",
            "5.  Access Weaviate at `http://localhost:8080` and configure as needed.\n",
            "6.  Check if Weaviate is ready by hitting the readiness endpoint:\n",
            "    ```bash\n",
            "    curl localhost:8080/v1/.well-known/ready\n",
            "    ```\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "prompt = \"\"\"\n",
        "You are connected to a database that has a blog post on deploying Weaviate on Docker.\n",
        "Can you answer how I can Weaviate with Docker?\n",
        "\"\"\"\n",
        "\n",
        "chat = client.chats.create(model='gemini-2.0-flash', config=config)\n",
        "response = chat.send_message(prompt)\n",
        "print(response.text)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "query-agent-as-a-tool.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
